import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn import svm
import joblib

# 1. 读取训练数据
train_data = pd.read_csv("任务三词语.csv")
texts = train_data["词语"]  # 输入：词语
labels = train_data["类别"] # 输出：类别

# 2. 将文本转换为数字（模型只能理解数字）
vectorizer = TfidfVectorizer()
X_train = vectorizer.fit_transform(texts)

# 3. 训练模型
#model = svm.SVC(probability=True)  # 启用概率估计
model = LogisticRegression(
    class_weight='balanced',  # 平衡类别权重
    multi_class='multinomial',  # 多分类问题
    solver='lbfgs',  # 优化算法
    max_iter=1000  # 增加迭代次数确保收敛
)
model.fit(X_train, labels)

# 4. 保存模型和转换器
joblib.dump(model, "../../hexinwenjianbao/text_classifier.pkl")      # 保存模型
joblib.dump(vectorizer, "../../hexinwenjianbao/tfidf_vectorizer.pkl")# 保存文本转数字的工具

print("模型训练完成！")
